//
// Created by Hoyin on 2024/4/13.
//

#ifndef VISUAL_LOCATION_DATASETLOADER_H
#define VISUAL_LOCATION_DATASETLOADER_H

#include <iostream>
#include <fstream>
#include <queue>
#include <filesystem>
#include <thread>
#include <chrono>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include "Clock.h"

namespace fs = std::filesystem;
template <class T>
class DatasetLoader {
protected:
	std::string _path;
	std::queue<T> _dataQueue;
	uint32_t _maxQueueSize{};
	uint32_t _dataSize{};
	uint32_t _outCount{};
	uint32_t _inCount{};
	bool _halt = false;

	std::vector<std::string> getFiles(const fs::path& dir);

public:
	explicit DatasetLoader(std::string path) : _path(std::move(path)) {};

	/**
	 * 加载数据到队列
	 */
	virtual void load() = 0;

	/**
	 * 中断信号
	 */
	inline void halt() {
		_halt = true;
	}

	/**
	 * 数据队列终了
	 * @return
	 */
	inline bool ended() {
		return _dataSize == 0 || _dataSize == _outCount;
	}

	/**
	 * 从队列头部取出一帧数据
	 * @return
	 */
	virtual T pop() = 0;

	inline size_t getOutCount() {
		return _outCount;
	}

	inline size_t getInCount() {
		return _inCount;
	}

	inline size_t getDataSize() {
		return _dataSize;
	}

	inline size_t getQueueSize() {
		return _dataQueue.size();
	}
};


class ImageLoader : public DatasetLoader<cv::Mat> {
private:
	std::vector<std::string> _files;
	void _handle();

public:
	ImageLoader(const std::string &dir, uint32_t queueSize) : DatasetLoader(dir) {
		_files = getFiles(dir);
		_dataSize = _files.size();
		_maxQueueSize = queueSize;
		_outCount = 0;
		_inCount = 0;
	}

	void load() override;

	cv::Mat pop() override;
};

template<class T>
std::vector<std::string> DatasetLoader<T>::getFiles(const fs::path &dir) {
	std::vector<std::string> file_paths;
	if (!fs::is_directory(dir)) return file_paths;
	for (const auto &entry: fs::directory_iterator(dir)) { // 遍历目录
		if (fs::is_directory(entry)) continue;
		file_paths.push_back(entry.path().string()); // 添加文件路径
	}
	return file_paths;
}

#endif //VISUAL_LOCATION_DATASETLOADER_H
